import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
import sys
sys.path.append("/home/project/Qtrainer_slim/test/yolov4") 
from configs import yolov4_cfg
from collections import OrderedDict

class YoloV4CspNet(nn.Module):

  
    def __init__(self):
        super(YoloV4CspNet,self).__init__()

        backboneConfig = yolov4_cfg.yolov4_cfg.configs
        # print("ddd",len(backboneConfig))
        
        self.layerList = []
        # self.yolov4backbone = nn.Sequential()
        self.yolov4Layer = nn.Sequential()

        for i in range(len(backboneConfig)):
            index = backboneConfig[i][0]
            in_channels= backboneConfig[i][1]
            out_channels= backboneConfig[i][2]
            kernel_size= backboneConfig[i][3]
            stride= backboneConfig[i][4]
            padding= backboneConfig[i][5]
            momentum= backboneConfig[i][6]
            eps= backboneConfig[i][7]
            # “conv”,"route","shortcut","upsample","yolohead","maxpool" 
            layertype = backboneConfig[i][8]
            layersp = backboneConfig[i][9]

            layername = str(index)
            # 卷积块
            if(layertype == "conv"):
                self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)) 
            # concate                 
            # elif(layertype == "route"):
            #     print("shortcut",layersp)
            #     routeLayer = {}
            #     for i in range(len(layersp)):
            #         routeName = "route_"+str(layersp[i])
            #         routeLayer[routeName] = self.yolov4Layer[layersp[i]]
            #         print("routeLayer", routeLayer)

            #     self.yolov4Layer.add_module(layername,torch.cat(**routeLayer))
            # # shortcut    
            # elif(layertype == "shortcut"):
            #     print("shortcut",layersp)
               
            #     self.yolov4Layer.add_module(layername,torch.add(self.yolov4Layer[layersp[0]],self.yolov4Layer[layersp[1]]))
            #         # torch.cat()

            #     # self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
            # # upsample    
            # elif(layertype == "upsample"):
            #     self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)) 
            # # yolohead    
            # elif(layertype == "yolohead"):
            #     self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
            # # maxpool     
            # elif(layertype == "maxpool"):
            #     self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)) 

            print("index",index)
            layername = index
            
            self.layerList.append(layername)

        # self.fc1 = nn.Linear(32*40*40,2)

    def forward(self,x):

        backboneConfig = yolov4_cfg.yolov4_cfg.configs
        # print("ddd",len(backboneConfig))
        for i in range(len(backboneConfig)):
            index = backboneConfig[i][0]
            in_channels= backboneConfig[i][1]
            out_channels= backboneConfig[i][2]
            kernel_size= backboneConfig[i][3]
            stride= backboneConfig[i][4]
            padding= backboneConfig[i][5]
            momentum= backboneConfig[i][6]
            eps= backboneConfig[i][7]
            # “conv”,"route","shortcut","upsample","yolohead","maxpool" 
            layertype = backboneConfig[i][8]
            layersp = backboneConfig[i][9]
            layername = str(index)
            # 每层输出值
            layerValueList = []
            # 卷积块
            if(layertype == "conv"):
                x = self.yolov4Layer[layername](x)
                layerValueList.append(x)
            # concate                 
            elif(layertype == "route"):
                print("shortcut",layersp)
                routeValueList = {}
                for i in range(len(layersp)):
                    routeName = "route_"+str(layersp[i])
                    routeValueList[routeName] = layerValueList[layersp[i]]
                    print("routeLayer", routeLayerValue)
                x = torch.cat(**routeValueList)
                
                layerValueList.append(x)

                
            # shortcut    
            elif(layertype == "shortcut"):
                print("shortcut",layersp)
                x = torch.add(self.yolov4Layer[layersp[0]](x),)
                self.yolov4Layer.add_module(layername,torch.add(self.yolov4Layer[layersp[0]],self.yolov4Layer[layersp[1]]))

                layerValueList.append(x)
                    # torch.cat()

                # self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
            # upsample    
            elif(layertype == "upsample"):
                self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps)) 
                layerValueList.append(x)
            # yolohead    
            elif(layertype == "yolohead"):
                self.yolov4Layer.add_module(layername, Convolutional(in_channels,out_channels,kernel_size,stride,padding,momentum,eps))
                layerValueList.append(x)
            # maxpool     
            elif(layertype == "maxpool"):
                x = F.max_pool2d()
                layerValueList.append(x)
            

        
        # x = self.yolov4Layer(x)
        # x = x.view(-1,self.num_flat_features(x))
        # x = F.relu(self.fc1(x))
        return x
        
     #使用num_flat_features函数计算张量x的总特征量（把每个数字都看出是一个特征，即特征总量），比如x是4*2*2的张量，那么它的特征总量就是16。
    def num_flat_features(self, x):
        size = x.size()[1:] # 这里为什么要使用[1:],是因为pytorch只接受批输入，也就是说一次性输入好几张图片，那么输入数据张量的维度自然上升到了4维
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

# yolo卷积块
class Convolutional(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
        super(Convolutional,self).__init__()

        self.convl = nn.Conv2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding)
        self.bn1 = nn.BatchNorm2d(out_channels, momentum = momentum, eps = eps, affine = True)

    def forward(self,x):
        x = self.convl(x)
        x = self.bn1(x) 
        x = F.relu(x)
        return x



# shortcut
class ShortCut(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
        super(ShortCut,self).__init__()

    def forward(self,x1,x2):
        x = torch.add(x1,x2)
        return x


# route 
class Route(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
        super(Route,self).__init__()

    def forward(self,tuple,dim=0):
        x = torch.cat(**tuple)
        return x


#  YoloHead 
class YoloHead(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,momentum,eps):
        super(YoloHead,self).__init__()

        self.convl = nn.Conv2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding)
        self.bn1 = nn.BatchNorm2d(out_channels, momentum = momentum, eps = eps, affine = True)

    def forward(self,x):
        x = self.convl(x)
        x = self.bn1(x) 
        x = F.relu(x)
        return x


 




        